from __future__ import print_function, absolute_import

import os.path as osp

import numpy as np
from scipy.io import loadmat


class VesselImageDataset(object):
    def __init__(self, root, min_seq_len=0):
        """
        @param root: 训练集所在路径
        @param min_seq_len: 单条轨迹的最少图片数，默认为0
        """
        self.root = root
        self.train_name_path = osp.join(self.root, 'info/train_name.txt')
        self.test_name_path = osp.join(self.root, 'info/test_name.txt')
        self.track_train_info_path = osp.join(self.root, 'info/tracks_train_info.mat')
        self.track_test_info_path = osp.join(self.root, 'info/tracks_test_info.mat')
        self.query_IDX_path = osp.join(self.root, 'info/query_IDX.mat')
        self._check_before_run()
        # 读取info文件
        train_names = self._get_names(self.train_name_path)
        test_names = self._get_names(self.test_name_path)
        track_train = loadmat(self.track_train_info_path)['tracks_train_info']  # numpy.ndarray (8298, 4)
        track_test = loadmat(self.track_test_info_path)['tracks_test_info']  # numpy.ndarray (12180, 4)
        query_IDX = loadmat(self.query_IDX_path)['query_IDX'].squeeze()  # numpy.ndarray (1980,)
        query_IDX -= 1  # index from 0，因为mat中的索引是从1开始的
        track_query = track_test[query_IDX, :]
        gallery_IDX = [i for i in range(track_test.shape[0]) if i not in query_IDX]
        track_gallery = track_test[gallery_IDX, :]
        train, num_train_tracklets, num_train_pids, num_train_imgs_per_tracklet = \
            self._process_data(names=train_names,
                               meta_data=track_train,
                               home_dir='images',
                               relabel=True,
                               min_seq_len=min_seq_len)
        query, num_query_tracklets, num_query_pids, num_query_imgs_per_tracklet = \
            self._process_data(names=test_names,
                               meta_data=track_query,
                               home_dir='images',
                               relabel=False,
                               min_seq_len=min_seq_len)
        gallery, num_gallery_tracklets, num_gallery_pids, num_gallery_imgs_per_tracklet = \
            self._process_data(names=test_names,
                               meta_data=track_gallery,
                               home_dir='images',
                               relabel=False,
                               min_seq_len=min_seq_len)

        num_imgs_per_tracklet = num_train_imgs_per_tracklet + num_query_imgs_per_tracklet + num_gallery_imgs_per_tracklet
        min_num = np.min(num_imgs_per_tracklet)
        max_num = np.max(num_imgs_per_tracklet)
        avg_num = np.mean(num_imgs_per_tracklet)

        num_total_pids = num_train_pids + num_query_pids
        num_total_tracklets = num_train_tracklets + num_query_tracklets + num_gallery_tracklets

        print("=> Dataset loaded")
        print("Dataset statistics:")
        print("  ------------------------------")
        print("  subset   | # ids | # tracklets")
        print("  ------------------------------")
        print("  train    | {:5d} | {:8d}".format(num_train_pids, num_train_tracklets))
        print("  query    | {:5d} | {:8d}".format(num_query_pids, num_query_tracklets))
        print("  gallery  | {:5d} | {:8d}".format(num_gallery_pids, num_gallery_tracklets))
        print("  ------------------------------")
        print("  total    | {:5d} | {:8d}".format(num_total_pids, num_total_tracklets))
        print("  number of images per tracklet: {} ~ {}, average {:.1f}".format(min_num, max_num, avg_num))
        print("  ------------------------------")

        self.train = train
        self.query = query
        self.gallery = gallery

        self.num_train_pids = num_train_pids
        self.num_query_pids = num_query_pids
        self.num_gallery_pids = num_gallery_pids

    def _check_before_run(self):
        """Check if all files are available before going deeper"""
        if not osp.exists(self.root):
            raise RuntimeError("'{}' is not available".format(self.root))
        if not osp.exists(self.train_name_path):
            raise RuntimeError("'{}' is not available".format(self.train_name_path))
        if not osp.exists(self.test_name_path):
            raise RuntimeError("'{}' is not available".format(self.test_name_path))
        if not osp.exists(self.track_train_info_path):
            raise RuntimeError("'{}' is not available".format(self.track_train_info_path))
        if not osp.exists(self.track_test_info_path):
            raise RuntimeError("'{}' is not available".format(self.track_test_info_path))
        if not osp.exists(self.query_IDX_path):
            raise RuntimeError("'{}' is not available".format(self.query_IDX_path))

    def _get_names(self, fpath):
        names = []
        with open(fpath, 'r') as f:
            for line in f:
                new_line = line.rstrip()
                names.append(new_line)
        return names

    def _process_data(self, names, meta_data, home_dir=None, relabel=False, min_seq_len=0):
        """
        @param names: 图片名构成的list
        @param meta_data: 各条轨迹的信息（轨迹起始图片名 轨迹结束图片名 该条轨迹船舶id名）
        @param home_dir: 图片所在路径
        @param relabel:
        @param min_seq_len:
        @return:
        """
        num_tracklets = meta_data.shape[0]
        pid_list = list(set(meta_data[:, 2].tolist()))
        num_pids = len(pid_list)
        # 标签重划分：从0开始构建id标签
        if relabel: pid2label = {pid: label for label, pid in enumerate(pid_list)}
        tracklets = []
        num_imgs_per_tracklet = []

        for tracklet_idx in range(num_tracklets):
            data = meta_data[tracklet_idx, ...]  # 单条轨迹的信息
            start_index, end_index, pid = data
            if pid == -1: continue  # junk images are just ignored
            img_names = names[start_index - 1:end_index]  # 提取该条轨迹的所有图片名

            if relabel: pid = pid2label[pid]

            # make sure image names correspond to the same vessel
            pnames = [img_name.split('_')[1] for img_name in img_names]
            assert len(set(pnames)) == 1, "Error: a single tracklet contains different person images"

            # append image names with directory information
            img_paths = [osp.join(self.root, home_dir, img_name) for img_name in img_names]
            if len(img_paths) >= min_seq_len:
                img_paths = tuple(img_paths)
                tracklets.append((img_paths, pid))
                num_imgs_per_tracklet.append(len(img_paths))

        num_tracklets = len(tracklets)
        return tracklets, num_tracklets, num_pids, num_imgs_per_tracklet
